"""
实时数据处理模块
"""
from typing import Dict, List, Optional
import pandas as pd
import ccxt
from datetime import datetime, timedelta

from .base import DataHandler
from utils.retry import retry_on_exception

class LiveDataHandler(DataHandler):
    """实时数据处理器"""
    
    def __init__(self, exchanges: Dict[str, ccxt.Exchange], symbols: List[str], timeframe: str = '1h'):
        """初始化实时数据处理器
        
        Args:
            exchanges: 交易所实例字典，键为交易所名称
            symbols: 交易对列表
            timeframe: 时间周期
        """
        super().__init__()
        self.exchanges = exchanges
        self.symbols = symbols
        self.timeframe = timeframe
        
        # 初始化数据结构
        self._init_data()
    
    def _init_data(self):
        """初始化数据结构"""
        for exchange_name in self.exchanges:
            if exchange_name not in self._data:
                self._data[exchange_name] = {}
                self._latest_data[exchange_name] = {}
                self._current_idx[exchange_name] = {}
            
            for symbol in self.symbols:
                self._data[exchange_name][symbol] = pd.DataFrame()
                self._latest_data[exchange_name][symbol] = []
                self._current_idx[exchange_name][symbol] = 0
    
    @retry_on_exception()
    def _fetch_ohlcv(self, exchange: ccxt.Exchange, symbol: str, **kwargs) -> List:
        """获取K线数据"""
        timeframe = kwargs.pop('timeframe', self.timeframe)
        limit = kwargs.pop('limit', 100)
        
        # 获取最近的K线数据
        ohlcv = exchange.fetch_ohlcv(
            symbol=symbol,
            timeframe=timeframe,
            limit=limit,
            **kwargs
        )
        return ohlcv
    
    def update_bars(self):
        """更新最新的K线数据"""
        for exchange_name, exchange in self.exchanges.items():
            for symbol in self.symbols:
                try:
                    # 获取最新数据
                    ohlcv = self._fetch_ohlcv(exchange, symbol)
                    
                    # 转换为DataFrame
                    df = pd.DataFrame(
                        ohlcv,
                        columns=['timestamp', 'open', 'high', 'low', 'close', 'volume']
                    )
                    df['timestamp'] = pd.to_datetime(df['timestamp'], unit='ms')
                    df.set_index('timestamp', inplace=True)
                    
                    # 更新数据
                    self._data[exchange_name][symbol] = pd.concat([
                        self._data[exchange_name][symbol],
                        df
                    ]).drop_duplicates()
                    
                    # 添加到最新数据列表
                    latest_bar = df.iloc[-1].to_dict()
                    latest_bar['timestamp'] = df.index[-1]
                    self._latest_data[exchange_name][symbol].append(latest_bar)
                    
                except Exception as e:
                    print(f"更新{exchange_name}的{symbol}数据时出错: {str(e)}")
                    continue
    
    def get_latest_bar(self, exchange: str, symbol: str) -> Optional[Dict]:
        """获取最新的一条数据"""
        try:
            latest_data = self._latest_data[exchange][symbol]
            return latest_data[-1] if latest_data else None
        except (KeyError, IndexError):
            return None
    
    def get_latest_bars(self, exchange: str, symbol: str, *, n: int = 50) -> List[Dict]:
        """获取最新的n条数据"""
        try:
            latest_data = self._latest_data[exchange][symbol]
            return latest_data[-n:] if len(latest_data) >= n else latest_data
        except (KeyError, IndexError):
            return []
